"""
Structed-Data Extraction for Images
"""

import numpy as np
class Tracking(object):
    def __init__(self):
        """
        self.object_detector = ObjectDetector()
        self.second_detector = SecondDetector()
        self.tracker = Sort(max_age = 3, min_hits = 5)#多路视频，可实例化多个追踪器
        """
        self.trajectory_store = {}#储存轨迹数据
        self.second_detect_freq = 3
    
    def update(self, frame):
        frame = np.transpose(frame, [2, 0, 1])#(H, W, C) --> (C, H, W)
        tracks = "[[306,323,306,323,'vehicle_type', 'vehicle_color', 'vehicle_plate', 'tracking_id'],[557,426,557,426,'vehicle_type', 'vehicle_color', 'vehicle_plate', 'tracking_id']]"
        """
        # start object_detection
        object_list = self.object_detector.detect(frame)
        #print(object_list)
        
        tracks = self.tracker.update(object_list, frame)
        
        #print('tracks: {}'.format(tracks))
        if len(tracks) > 0:
            #obj = [x1,y1,x2,y2, class_id, vehicle_color, plate_num, obj_id]
            
            ##判断是否要进行二次识别
            second_detect_index = []
            obj_box = []
            for i_index, obj in enumerate(tracks):
                #print(obj)
                if obj[4] not in vehicle_id:
                    continue
                if obj[-1] not in list(self.trajectory_store.keys()):
                    second_detect_index.append(i_index)
                    obj_box.append(obj)
                else:
                    temp = self.trajectory_store[obj[-1]]
                    if (temp[-1]+1)%self.second_detect_freq == 0:
                        second_detect_index.append(i_index)
                        obj_box.append(obj)
            #二次特征检测
            if len(obj_box) > 0:
                color_list,plate_list = self.second_detector.detect(frame, obj_box, frame_index)
                for i_index,color in enumerate(color_list):
                    tracks[second_detect_index[i_index]][5] = color
                    tracks[second_detect_index[i_index]][6] = plate_list[i_index]
                    
                
            #针对二次检测结果对追踪目标的特征进行修正，更新历史轨迹
            for obj in tracks:
                #print('obj: {}'.format(obj)) 
                
                if obj[-1] not in list(self.trajectory_store.keys()):
                    #trajectory_store[obj[-1]] = [obj +[frame_index]]
                    self.trajectory_store[obj[-1]] = [obj, 1]
                else:
                    temp = self.trajectory_store[obj[-1]]
                    track_count = temp[-1]
                    if (track_count+1)%self.second_detect_freq == 0:
                        self.trajectory_store[obj[-1]] = [obj, track_count+1]
                    else:
                        obj[5:] = temp[0][5:]
                        self.trajectory_store[obj[-1]] = [obj, track_count+1]
        """
        return tracks

